import time
import cv2
from stable_baselines3 import DQN
from stable_baselines3.common.vec_env import DummyVecEnv
from RL_dino import DinoGameEnv

# 加载模型并进行预测
env = DummyVecEnv([lambda: DinoGameEnv()])
model = DQN.load('D:/zhulibo/university/大三/智能系统控制/chrome-dinosaur/PyMini-Dino/train/best_model_35000')

for episode in range(5): 
    done = False
    total_reward = 0
    action, _ = model.predict(obs)
    obs, reward, done, info = env.step(int(action))
    time.sleep(0.01)
    total_reward += reward
    print('Total Reward for episode {} is {}'.format(episode, total_reward))
    time.sleep(2)